from numpy import linspace, pi, cos, sin


def curve(n_mer, n_long):
    phi = linspace(0, 2 * pi, 2000)
    return [cos(phi * n_mer) * (1 + 0.5 * cos(n_long * phi)),
            sin(phi * n_mer) * (1 + 0.5 * cos(n_long * phi)),
            0.5 * sin(n_long * phi),
            sin(phi * n_mer)]


from traits.api import *
from traitsui.api import *
from tvtk.pyface.scene_editor import SceneEditor
from mayavi.tools.mlab_scene_model import MlabSceneModel
from mayavi.core.ui.mayavi_scene import MayaviScene


class Visualization(HasTraits):
    meridional = Range(1, 30, 6)
    transverse = Range(0, 30, 11)
    scene = Instance(MlabSceneModel, ())
    button=Button(label='test')
    def _button_fired(self):
        print('fdjf')

    def __init__(self):
        # Do not forget to call the parent's __init__
        HasTraits.__init__(self)
        x1, y, z, t = curve(self.meridional, self.transverse)
        self.plot = self.scene.mlab.plot3d(x1, y, z, t, colormap='Spectral',figure=self.scene.mayavi_scene)

    @observe('meridional,transverse')
    def update_plot(self, event=None):
        x, y, z, t = curve(self.meridional, self.transverse)
        self.plot.mlab_source.trait_set(x=x, y=y, z=z, scalars=t)

    # the layout of the dialog created
    view = View(
        Item(
            'scene', editor=SceneEditor(scene_class=MayaviScene),
            height=250, width=300, show_label=False),
        HGroup(
            '_',
            'meridional',
            'transverse',
            UItem('button'),
        ),
    )
    def test(self):
        print('this is test in mlab')



form = Visualization()

if __name__ == '__main__':
    form.configure_traits()

